// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                       avtADIOS2BaseFileFormat.C                           //
// ************************************************************************* //

#include <avtMTSDFileFormatInterface.h>
#include <avtADIOS2BaseFileFormat.h>
#include <ADIOS2HelperFuncs.h>

#include <string>
#include <map>

#include <vtkFloatArray.h>
#include <vtkDoubleArray.h>
#include <vtkRectilinearGrid.h>
#include <vtkStructuredGrid.h>
#include <vtkUnstructuredGrid.h>

#include <avtDatabaseMetaData.h>

#include <Expression.h>

#include <InvalidVariableException.h>
#include <VisItStreamUtil.h>
#include <DebugStream.h>

using namespace std;

static std::vector<int> CSVToVectorInt(const std::string csv) noexcept
{
    std::vector<int> numbers;
    if (csv.empty())
    {
        return numbers;
    }

    string tmp;
    for (int i = 0; i < csv.size(); i++)
        if (csv[i] != ' ') tmp.push_back(csv[i]);

    if (tmp.find(",") == tmp.npos) // if no commas, one int
    {
        numbers.push_back(std::stoi(tmp)); // might need to be checked
    }
    else
    {
        numbers.reserve(std::count(tmp.begin(), tmp.end(), ','));
        std::istringstream csvSS(tmp);
        std::string value;
        while (std::getline(csvSS, value, ','))
        {
            numbers.push_back(std::stoi(value));
        }
    }

    return numbers;
}

bool
avtADIOS2BaseFileFormat::Identify(const char *fname)
{
    return true;
}

avtFileFormatInterface *
avtADIOS2BaseFileFormat::CreateInterface(const char *const *list,
                                         int nList,
                                         int nBlock)
{
    int nTimestepGroups = nList / nBlock;
    avtMTSDFileFormat ***ffl = new avtMTSDFileFormat**[nTimestepGroups];
    for (int i = 0; i < nTimestepGroups; i++)
    {
        ffl[i] =  new avtMTSDFileFormat*[nBlock];
        for (int j = 0; j < nBlock; j++)
            ffl[i][j] =  new avtADIOS2BaseFileFormat(list[i*nBlock +j]);
    }
    return new avtMTSDFileFormatInterface(ffl, nTimestepGroups, nBlock);
}

avtFileFormatInterface *
avtADIOS2BaseFileFormat::CreateInterfaceADIOS2(
        const char *const *list,
        int nList,
        int nBlock,
        std::shared_ptr<adios2::ADIOS> adios,
        adios2::Engine &reader,
        adios2::IO &io,
        std::map<std::string, adios2::Params> &variables,
        std::map<std::string, adios2::Params> &attributes
        )
{
    int nTimestepGroups = nList / nBlock;
    avtMTSDFileFormat ***ffl = new avtMTSDFileFormat**[nTimestepGroups];
    for (int i = 0; i < nTimestepGroups; i++)
    {
        ffl[i] =  new avtMTSDFileFormat*[nBlock];
        for (int j = 0; j < nBlock; j++)
        {
        debug5 << "----------- ADIOS Base create interface for  "
             << list[i*nBlock +j]
             << " -----------------" << endl;

            if (!i && !j)
            {
                ffl[i][j] =  new avtADIOS2BaseFileFormat(adios, reader, io, variables, attributes, list[i*nBlock +j]);
            }
            else
            {
                ffl[i][j] =  new avtADIOS2BaseFileFormat(list[i*nBlock +j]);
            }
        }
    }
    debug5 << "----------- ADIOS Base return MTSD interface " << endl;
    return new avtMTSDFileFormatInterface(ffl, nTimestepGroups, nBlock);
}


// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat constructor
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
//  Modifications:
//    Justin Privitera, Thu Jan 18 09:56:51 PST 2024
//    Removed adios2::DebugON since it is not present in newer versions of
//    adios2.
//
// ****************************************************************************


avtADIOS2BaseFileFormat::avtADIOS2BaseFileFormat(const char *filename)
    : adios(std::make_shared<adios2::ADIOS>()),
      numTimeSteps(1),
      isClosed(false),
      avtMTSDFileFormat(&filename, 1)
{
    engineName = ADIOS2Helper_GetEngineName(filename);
    stagingMode = ADIOS2Helper_IsStagingEngine(engineName);
    string adiosFileName = ADIOS2Helper_GetFileName(filename);
    io = adios2::IO(adios->DeclareIO(engineName));
    io.SetEngine(engineName);
    reader = io.Open(adiosFileName, adios2::Mode::Read);

    if (stagingMode)
    {
        numTimeSteps = 100000;
    }
    else
    {
        map<string, adios2::Params> vars = io.AvailableVariables();
        for (auto &v : vars)
        {
            string nsteps = v.second["AvailableStepsCount"];
            if (!nsteps.empty())
            {
                numTimeSteps = std::stoi(nsteps);
                break;
            }
        }
    }
}

avtADIOS2BaseFileFormat::avtADIOS2BaseFileFormat(std::shared_ptr<adios2::ADIOS> adios,
        adios2::Engine &reader,
        adios2::IO &io,
        std::map<std::string, adios2::Params> &variables,
        std::map<std::string, adios2::Params> &attributes,
        const char *filename)
    : adios(adios),
      numTimeSteps(1),
      isClosed(false),
      reader(reader),
      io(io),
      variables(variables),
      avtMTSDFileFormat(&filename, 1)
{
    engineName = ADIOS2Helper_GetEngineName(filename);
    stagingMode = ADIOS2Helper_IsStagingEngine(engineName);

    if (stagingMode)
    {
        numTimeSteps = 100000;
    }
    else
    {
        for (auto &v : variables)
        {
            string nsteps = v.second["AvailableStepsCount"];
            if (!nsteps.empty())
            {
                numTimeSteps = std::stoi(nsteps);
                break;
            }
        }
    }
    debug5 << "----------- ADIOS Base reader for "
         << filename << " with " << numTimeSteps
         << "steps -----------------" << endl;
}

avtADIOS2BaseFileFormat::~avtADIOS2BaseFileFormat()
{
#if MDSERVER
    debug5 << "-----------\nADIOS destructor called in MDSERVER \n-----------------\n" << endl;
#else
    debug5 << "-----------\nADIOS destructor called in ENGINE \n-----------------\n" << endl;
#endif
    if (!isClosed)
    {
        debug5 << "-----------\nADIOS reader closes down stream \n-----------------\n" << endl;
        reader.Close();
    }
    isClosed = true;
}

// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat::GetNTimesteps
//
//  Purpose:
//      Tells the rest of the code how many timesteps there are in this file.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

int
avtADIOS2BaseFileFormat::GetNTimesteps()
{
    return numTimeSteps;
}


// ****************************************************************************
//  Method: avtADIOS2FileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

void
avtADIOS2BaseFileFormat::FreeUpResources(void)
{
}


// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

void
avtADIOS2BaseFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md, int timeState)
{
    if (stagingMode)
    {
        /*
        adios2::StepStatus status = reader.BeginStep(adios2::StepMode::Read, -1.0f);
        if (status != adios2::StepStatus::OK)
            return;
#if MDSERVER
        debug5<<" MDSERVER received streaming step = "<<reader.CurrentStep()<<endl;
#else
        debug5<<" Server populates metadata from streaming step = "<<reader.CurrentStep()<<endl;
#endif
        */
    }

    meshInfo.clear();
    variables = io.AvailableVariables();

    debug5<<" Server populates metadata with "<<variables.size() <<" variables."<<endl;

#if MDSERVER
    if (stagingMode)
    {
        reader.EndStep();
    }
#endif

    vector<pair<string,string>> vars;

    for (const auto &varInfo : variables)
    {
        auto params = varInfo.second;

        string shape = params["Shape"];
        auto vecShape = CSVToVectorInt(shape);
        if (meshInfo.find(shape) == meshInfo.end())
        {
            string meshNm;
            int dim = 0;
            if (vecShape.size() == 2)
            {
                meshNm = "mesh" + std::to_string(vecShape[0])+"x"+std::to_string(vecShape[1]);
                dim = 2;
            }
            else if (vecShape.size() == 3)
            {
                meshNm = "mesh" + std::to_string(vecShape[0])+"x"+std::to_string(vecShape[1])+"x"+std::to_string(vecShape[2]);
                dim = 3;
            }
            meshInfo[shape] = std::make_pair(dim, meshNm);
        }
        vars.push_back(std::make_pair(varInfo.first, meshInfo[shape].second));
    }

    for (auto &m : meshInfo)
    {
        string nm = m.second.second;
        int dim = m.second.first;
        AddMeshToMetaData(md, nm, AVT_RECTILINEAR_MESH, NULL, 1, 0, dim, dim);
    }

    for (auto &v : vars)
        AddScalarVarToMetaData(md, v.first, v.second, AVT_NODECENT);

#if MDSERVER
    if (!isClosed)
    {
        debug5 << "-----------\nADIOS reader closes down stream in MDSERVER \n-----------------\n" << endl;
        reader.Close();
    }
    isClosed = true;
#endif
}


// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      timestate   The index of the timestate.  If GetNTimesteps returned
//                  'N' time steps, this is guaranteed to be between 0 and N-1.
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataSet *
avtADIOS2BaseFileFormat::GetMesh(int timestate, const char *meshname)
{
    //Find the mesh info we need.
    string shapeStr;
    for (auto &m : meshInfo)
    {
        if (m.second.second == string(meshname))
        {
            shapeStr = m.first;
            break;
        }
    }
    if (shapeStr.size() == 0)
        return NULL;

    auto shape = CSVToVectorInt(shapeStr);
    if (shape.size() == 2)
        shape.push_back(1);

    //debug5<<"********** SWAP DIMS"<<endl;
    std::swap(shape[0], shape[1]);

    int dims[3] = {shape[0], shape[1], shape[2]};

    vtkRectilinearGrid *grid = vtkRectilinearGrid::New();
    grid->SetDimensions(dims);
    vtkFloatArray *coords[3] = {vtkFloatArray::New(),vtkFloatArray::New(),vtkFloatArray::New()};

    for (int c = 0; c < 3; c++)
    {
        coords[c]->SetNumberOfTuples(shape[c]);
        for (int i = 0; i < shape[c]; i++)
            coords[c]->SetTuple1(i, i);
    }

    grid->SetXCoordinates(coords[0]);
    grid->SetYCoordinates(coords[1]);
    grid->SetZCoordinates(coords[2]);
    coords[0]->Delete();
    coords[1]->Delete();
    coords[2]->Delete();

    return grid;
}


// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataArray *
avtADIOS2BaseFileFormat::GetVar(int timestate, const char *varname)
{
    static int prev_timestate = 0;
    debug5<<"GetVar: "<<varname<<" ts= "<<timestate<<endl;

    if (variables.find(varname) == variables.end())
        return NULL;

    auto var = variables[varname];
    string varType = var["Type"];

    if (stagingMode)
    {
        if (timestate != prev_timestate)
        {
            debug5<<" release step = "<<reader.CurrentStep()<<endl;
            reader.EndStep(); // release previous step
            prev_timestate = timestate;
            adios2::StepStatus status = adios2::StepStatus::NotReady;
            while (status == adios2::StepStatus::NotReady)
            {
                status = reader.BeginStep(adios2::StepMode::Read, 10.0f);
                if (status == adios2::StepStatus::EndOfStream ||
                    status == adios2::StepStatus::OtherError)
                {
                    reader.Close();
                    //variables.clear();
                    return NULL;
                }
                else if (status == adios2::StepStatus::NotReady)
                {
                    debug5<<" still waiting for next step = "
                        <<reader.CurrentStep()<<endl;
                }
            }
            debug5<<" received streaming step = "<<reader.CurrentStep()<<endl;
        }
    }


    adios2::Variable<double> v = io.InquireVariable<double>(varname);
    //debug5<<"DIMS= "<<v.Shape()<<endl;
    if (!stagingMode)
        v.SetStepSelection({timestate, 1});

    size_t numVals = 1;
    for (int i = 0; i < v.Shape().size(); i++)
        numVals *= v.Shape()[i];

    vector<double> buff(numVals);
    reader.Get(v, buff.data(), adios2::Mode::Sync);


    vtkDoubleArray *arr = vtkDoubleArray::New();
    arr->SetNumberOfComponents(1);
    arr->SetNumberOfTuples(numVals);
    for (int i = 0; i < numVals; i++)
        arr->SetTuple1(i, buff[i]);

    return arr;
}


// ****************************************************************************
//  Method: avtADIOS2BaseFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataArray *
avtADIOS2BaseFileFormat::GetVectorVar(int timestate, const char *varname)
{
    return NULL;
}
